#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/mman.h>
#include <sys/types.h>

#define PHY_ADDR  0x38000000
#define MAP_SIZE  0x40000

/* control register */
#define ACTIVE_REG 0
#define DONE_REG 1
#define START_REG 2
#define END_REG 3

struct trandata {
  unsigned int number;
  unsigned int data;
};

unsigned char* load_png(const char *filename, size_t *size);
int write_to_mem(unsigned char *data, size_t size);

// global variables
static int fd;
static void *map_base = NULL;
static FILE *fp;
static struct trandata buf;

void write_reg(int fd, unsigned int number, unsigned int data)
{
  buf.number = number;
  buf.data = data;
  if (write(fd, &buf, sizeof(buf)) < 0) {
    printf("Write reg %d failed!\r\n",number);
    munmap(map_base, MAP_SIZE);
    close(fd);
    exit(1);
  }
}

void read_reg(int fd, unsigned int number, unsigned int * data)
{
  buf.number = number;
  buf.data = 0;
  read(fd, &buf, sizeof(buf));
  *data = buf.data;
}

int main(int argc, char *argv[])
{
  struct timeval start, end;
  unsigned int done_data = 0, start_data = 0, end_data = 0;
  unsigned long time;

  if (argc < 3) {
    printf("Usage: %s <file> <out_file>\r\n", argv[0]);
    return -1;
  }

  // 1. load image data to virtual memory
  const char *filename = argv[1];
  size_t png_size;

  unsigned char *png_data = load_png(filename, &png_size);
  if (!png_data) {
    fprintf(stderr, "Failed to load PNG data from %s\r\n", filename);
    return EXIT_FAILURE;
  }
  printf("Loaded %zu bytes from %s\r\n", png_size, filename);


  // 2. copy image data to 0x38000000
  // 2.1 open "/dev/mem"
  fd = open("/dev/mem", O_RDWR | O_SYNC);
  if (fd < 0) {
    perror("open /dev/mem");
    return -1;
  }

  // 2.2 map physical address 0x38000000 to user virtual address
  map_base = mmap(NULL, MAP_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, fd, (off_t)PHY_ADDR);
  if (map_base == MAP_FAILED) {
    perror("mmap");
    close(fd);
    return -1;
  }
  printf("Physical address 0x%lX mapped at virtual address %p\r\n",
         (unsigned long)PHY_ADDR, map_base);

  // 2.3 transfer image data to this physical address
  memcpy(map_base, png_data, png_size);

  // Clean up
  close(fd);
  free(png_data);

  // 3. make a stop
  // printf("Please trigger hardware and input a character then\n");
  // getchar();

  // 3. trigger errdiff system
  // 3.1 open errdiff device
  fd = open("/dev/errdiff", O_RDWR);
  if (fd < 0) {
    printf("file /dev/errdiff open failed!\r\n");
    munmap(map_base, MAP_SIZE);
    return -1;
  }
  // 3.2 send active signal to errdiff
  write_reg(fd, ACTIVE_REG, 0x1);

  gettimeofday(&start, NULL);
  // while (1) { // load image to bram
  //   read_reg(fd, START_REG, &start_data);
  //   if (start_data) break;
  // }
  // while (1) { // the real process
  //   read_reg(fd, END_REG, &end_data);
  //   if (end_data) break;
  // }
  while (1) {
    read_reg(fd, DONE_REG, &done_data);
    if (done_data) break;
  }
  gettimeofday(&end, NULL);
  printf("image process finished!\r\n");
  close(fd);

  // 4. save process image data
  const char *outfile = argv[2];
  fp = fopen(outfile, "wb");
  if (!fp) {
    perror("fopen");
    munmap(map_base, MAP_SIZE);
    return -1;
  }

  // 4.1 write processed image data to local file
  size_t written = fwrite(map_base, 1, MAP_SIZE, fp);
  if (written != MAP_SIZE) {
    fprintf(stderr, "fwrite: Written %zu bytes, expected %d bytes\n", written, MAP_SIZE);
    fclose(fp);
    munmap(map_base, MAP_SIZE);
    return -1;
  }
  printf("Processed data has been written!\r\n");

  // Clean up
  fclose(fp);
  munmap(map_base, MAP_SIZE);
  time = 1000000 * (end.tv_sec-start.tv_sec) + end.tv_usec-start.tv_usec;
  printf("total time = %ld us\r\n",time);

  return 0;
}

/**
 * @brief 从指定的 PNG 文件中读取二进制数据到内存缓冲区
 *
 * @param filename PNG 文件路径
 * @param size    输出参数，文件大小
 * @return unsigned char* 成功时返回指向数据的指针，失败时返回 NULL
 */
unsigned char* load_png(const char *filename, size_t *size)
{
    int fd = open(filename, O_RDONLY);
    if (fd < 0) {
        perror("open");
        return NULL;
    }

    struct stat st;
    if (fstat(fd, &st) < 0) {
        perror("fstat");
        close(fd);
        return NULL;
    }
    size_t filesize = st.st_size;

    unsigned char *buffer = (unsigned char *)malloc(filesize);
    if (!buffer) {
        perror("malloc");
        close(fd);
        return NULL;
    }

    ssize_t total_read = 0;
    while (total_read < filesize) {
        ssize_t ret = read(fd, buffer + total_read, filesize - total_read);
        if (ret < 0) {
            perror("read");
            free(buffer);
            close(fd);
            return NULL;
        }
        if (ret == 0) {
            break; // 文件结束
        }
        total_read += ret;
    }

    close(fd);

    if (total_read != filesize) {
        fprintf(stderr, "Incomplete read: expected %zu, got %zd\n", filesize, total_read);
        free(buffer);
        return NULL;
    }

    *size = filesize;
    return buffer;
}


/**
 * @brief 将数据写入通过 /dev/mem 映射的物理地址
 *
 * @param data 要写入的数据
 * @param size 数据大小（字节）
 * @return int 成功返回 0，失败返回 -1
 */
int write_to_mem(unsigned char *data, size_t size)
{
    int fd = open("/dev/mem", O_RDWR | O_SYNC);
    if (fd < 0) {
        perror("open /dev/mem");
        return -1;
    }

    /* 确保物理地址对齐到页面大小 */
    off_t target = PHY_ADDR;
    size_t aligned_size = ((size + (sysconf(_SC_PAGESIZE) - 1)) / sysconf(_SC_PAGESIZE)) * sysconf(_SC_PAGESIZE);

    void *mapped_base = mmap(NULL, aligned_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, target);
    if (mapped_base == MAP_FAILED) {
        perror("mmap");
        close(fd);
        return -1;
    }

    /* 使用 memcpy 将数据复制到映射的内存区域 */
    memcpy(mapped_base, data, size);

    /* 处理可能需要的缓存刷新 */
    // 如果硬件需要缓存一致性，可以调用 msync()
    if (msync(mapped_base, aligned_size, MS_SYNC) < 0) {
        perror("msync");
        munmap(mapped_base, aligned_size);
        close(fd);
        return -1;
    }

    /* 解除映射并关闭文件描述符 */
    if (munmap(mapped_base, aligned_size) == -1) {
        perror("munmap");
        close(fd);
        return -1;
    }

    close(fd);
    return 0;
}

